#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Retcode Statistics Analysis Tool
Used to analyze retcode distribution during Locust load testing
"""

import json
import time
from collections import defaultdict
from datetime import datetime
import csv
import os

class RetcodeStatsCollector:
    """Retcode Statistics Collector"""
    
    def __init__(self):
        self.stats = defaultdict(int)
        self.start_time = time.time()
        self.total_requests = 0
        
    def record_retcode(self, retcode):
        """Record retcode"""
        self.stats[retcode] += 1
        self.total_requests += 1
        
    def get_stats(self):
        """Get statistics information"""
        return dict(self.stats)
        
    def get_success_rate(self):
        """Get success rate"""
        if self.total_requests == 0:
            return 0.0
        success_count = self.stats.get(0, 0)
        return (success_count / self.total_requests) * 100
        
    def get_error_rate(self):
        """Get error rate"""
        return 100 - self.get_success_rate()
        
    def print_summary(self):
        """Print statistics summary"""
        print("\n" + "="*50)
        print("Retcode Statistics Summary")
        print("="*50)
        print(f"Test Start Time: {datetime.fromtimestamp(self.start_time)}")
        print(f"Total Requests: {self.total_requests}")
        print(f"Success Rate: {self.get_success_rate():.2f}%")
        print(f"Error Rate: {self.get_error_rate():.2f}%")
        print("\nRetcode Distribution:")
        
        # Display sorted by retcode
        sorted_stats = sorted(self.stats.items())
        for retcode, count in sorted_stats:
            percentage = (count / self.total_requests) * 100 if self.total_requests > 0 else 0
            status = "Success" if retcode == 0 else "Failure"
            print(f"  retcode {retcode} ({status}): {count} times ({percentage:.2f}%)")
            
    def export_to_csv(self, filename=None):
        """Export statistics data to CSV file"""
        if filename is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"retcode_stats_{timestamp}.csv"
            
        with open(filename, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(['Retcode', 'Count', 'Percentage', 'Status'])
            
            sorted_stats = sorted(self.stats.items())
            for retcode, count in sorted_stats:
                percentage = (count / self.total_requests) * 100 if self.total_requests > 0 else 0
                status = "Success" if retcode == 0 else "Failure"
                writer.writerow([retcode, count, f"{percentage:.2f}%", status])
                
        print(f"\nStatistics data exported to: {filename}")
        
    def export_to_json(self, filename=None):
        """Export statistics data to JSON file"""
        if filename is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"retcode_stats_{timestamp}.json"
            
        data = {
            'start_time': self.start_time,
            'total_requests': self.total_requests,
            'success_rate': self.get_success_rate(),
            'error_rate': self.get_error_rate(),
            'retcode_stats': dict(self.stats),
            'export_time': time.time()
        }
        
        with open(filename, 'w', encoding='utf-8') as jsonfile:
            json.dump(data, jsonfile, indent=2, ensure_ascii=False)
            
        print(f"Statistics data exported to: {filename}")

# Global statistics collector instance
stats_collector = RetcodeStatsCollector()

def record_retcode(retcode):
    """Convenient function to record retcode"""
    stats_collector.record_retcode(retcode)
    
def print_stats():
    """Convenient function to print statistics"""
    stats_collector.print_summary()
    
def export_stats(format='csv', filename=None):
    """Convenient function to export statistics"""
    if format.lower() == 'csv':
        stats_collector.export_to_csv(filename)
    elif format.lower() == 'json':
        stats_collector.export_to_json(filename)
    else:
        print(f"Unsupported format: {format}")

if __name__ == "__main__":
    # Example usage
    print("Retcode Statistics Tool Example")
    
    # Simulate some data
    import random
    for _ in range(1000):
        # Simulate different retcodes
        retcode = random.choices([0, 1, 2, 3, 'http_error'], weights=[80, 10, 5, 3, 2])[0]
        record_retcode(retcode)
    
    # Print statistics
    print_stats()
    
    # Export data
    export_stats('csv')
    export_stats('json')